#include "ntuples.C"
#include "PlotQCD.C"
#include <stdlib.h> 

/**
  mode: 0 for data, else for mc
**/

TString mycut(const char* sample, const char* jetNumberCriteria=">=1", int doTight=0) {

    TString cut_jet      = "GoodJet25_N " + TString(jetNumberCriteria); // number criteria for jet > 25GeV
    TString cut_lepton   = "GoodMu_N + GoodEl_N < 2"; // no more than 1 tight lepton 
    TString cut_trig_mc  = "EF_mu10_MSonly";
    TString cut_trig_dat = "((RunNumber<=160879 && L1_MU10 > 0 ) ||(RunNumber >= 160899 && RunNumber <= 162882 && EF_mu10_MSonly ) || (RunNumber >= 165591 && RunNumber <= 166143 && EF_mu13 ) || (RunNumber >= 166198 && EF_mu13_tight) )";            // muon trigger
    TString cut_etmiss   = "EtMissCut";         // Etmiss > 20 GeV
    TString cut_vertex   = "NonColBkgCut";      // having primary vertex
    TString cut_bad      = "CleaningCut";       // no bad jet
    TString cut_grl      = "grlevent";          // GRL events 
                         
    TString cut_tria     = "TriangularCut";     // triangular cut
    TString cut_btag     = "BtagCut";           // btagging
    //TString cut_other    = "emuOverlapCut && Jet25Cut";
    TString cut_other    = "emuOverlapCut";

    TString AND = " && ";
    TString OR  = " || ";

    TString cut = cut_jet    + AND + 
                  cut_lepton + AND + 
                  cut_etmiss + AND + 
                  cut_vertex + AND + 
    //              cut_tria   + AND +
                  cut_btag   + AND +
                  cut_other  + AND + 
                  cut_bad;

    // specific trigger cut according sample type 
    if ( TString(sample).BeginsWith("data10") ) {
        cut = cut + AND + cut_grl + AND + cut_trig_dat;
    }

    if ( TString(sample).BeginsWith("mc09") ) {
        cut = cut + AND + cut_trig_mc;
    }

    // tigger cut includs: triangular cut + btagging cut
    if ( doTight == 1 ) {
        //cut = cut + AND + "GoodMuonCut";
        //cut = cut + AND + "GoodMuonCut";
        cut = cut + AND + cut_tria;
    }

    cout << "event cut: " << cut << endl;

    return cut;
}

void eff_loose2tight(const char* sample, const char* jetNumberCriteria=">=1") {

    TH1D *h_eff    = new TH1D("h_eff", "h_eff", 100, 0., 25.);

    TString event_cut = mycut(sample, jetNumberCriteria);

    TString loose_cut = event_cut + " && GoodMuLoose_dR_jet > 0.4";
    TString tight_cut = event_cut + " && GoodMuLoose_dR_jet > 0.4 && GoodMuLoose_etcone30 < 4000. && GoodMuLoose_ptcone30 < 4000.";

    cout << "loose cut: " << loose_cut << endl;
    cout << "tight cut: " << tight_cut << endl;

    TH1D *h_loose  = load_histogram(sample, "loose", "GoodMuLoose_d0_sig_exPV", loose_cut, 100,  0., 25.);
    TH1D *h_tight  = load_histogram(sample, "tight", "GoodMuLoose_d0_sig_exPV", tight_cut, 100,  0., 25.);

    for (Int_t i=0; i < h_loose->GetNbinsX() + 1; i++) {
        Double_t N_loose = h_loose->GetBinContent(i);
        Double_t N_tight = h_tight->GetBinContent(i);
        Double_t bin_val = h_loose->GetBinCenter(i);

        Double_t eff = ( N_loose <= 0.) ? 0.:N_tight / N_loose;

        // calculating error of loose-to-tight efficiency as a function of (N_tight, dN=N_loose-N_tight)
        Double_t dN     = N_loose - N_tight;
        Double_t dN_sig = TMath::Sqrt(dN);
        //Double_t Nt_sig = TMath::Sqrt(N_tight);

        Double_t stat_err = 0.0;

        if ( eff != 0. ) {
            stat_err = TMath::Sqrt( 
                              TMath::Power( (N_loose + N_tight) / TMath::Power(N_loose,2) , 2) * N_tight +
                              TMath::Power( N_tight / TMath::Power(N_loose,2), 2) * dN
                       );
        }

        h_eff->SetBinContent(i,eff);
        h_eff->SetBinError(i,stat_err);

        cout << bin_val << " " << N_tight << " " << N_loose << " " << eff << " " << stat_err << endl;
    }

    // writing out the h_eff
    TFile *rsult = new TFile("eff_loose2tight.root", "recreate");

    h_loose->Write();
    h_tight->Write();
    h_eff->Write();
    rsult->Close();

    gDirectory->Clear();
}

void cal_eff_fake(double d0_sig, const char* sample="data10", const char* jetNumberCriteria=">=1") {

    TTree *myTree = load_sample(sample);

    TString d0_sig_thr = ""; d0_sig_thr +=  d0_sig;

    TString event_cut = mycut(sample,jetNumberCriteria);

    TString loose_cut = event_cut + " && GoodMuLoose_dR_jet > 0.4";
    TString tight_cut = loose_cut + " && GoodMuLoose_etcone30 < 4000. && GoodMuLoose_ptcone30 < 4000.";

    myTree->Project("h_loose", "GoodMuLoose_pt", loose_cut + " && GoodMuLoose_d0_sig_exPV > " + d0_sig_thr);
    myTree->Project("h_tight", "GoodMuLoose_pt", tight_cut + " && GoodMuLoose_d0_sig_exPV > " + d0_sig_thr);

    Double_t N_loose = h_loose->GetEntries();
    Double_t N_tight = h_tight->GetEntries();

    Double_t eff = ( N_loose <= 0.) ? 0.:N_tight / N_loose;

    // calculating error of loose-to-tight efficiency as a function of (N_tight, dN=N_loose-N_tight)
    Double_t dN     = N_loose - N_tight;
    Double_t dN_sig = TMath::Sqrt(dN);
    //Double_t Nt_sig = TMath::Sqrt(N_tight);

    Double_t stat_err = 0.0;

    if ( eff != 0. ) {
        stat_err = TMath::Sqrt( 
                          TMath::Power( (N_loose + N_tight) / TMath::Power(N_loose,2) , 2) * N_tight +
                          TMath::Power( N_tight / TMath::Power(N_loose,2), 2) * dN
                   );
    }

    cout << eff << " " << stat_err << endl;

    gDirectory->Clear();
}

void plot_observable(TString observable, TString prefix_hist, Int_t nbins, Double_t xlow, Double_t xhigh, const char* sample="data10", const char* jetNumberCriteria=">=1") {

    TString cut;
    TString cut_loose = " && GoodMuLoose_dR_jet > 0.4 ";
    TString cut_tight = cut_loose + " && GoodMuLoose_etcone30 < 4000. && GoodMuLoose_ptcone30 < 4000. ";

    // data
    cut = mycut(sample, jetNumberCriteria) + cut_tight;
    TH1D *h_data = load_histogram(sample, prefix_hist, observable, cut, nbins, xlow, xhigh);
    h_data->SetNameTitle(prefix_hist + "data", prefix_hist + "data");

    // ttbar
    cut = mycut("mc09_ttbar", jetNumberCriteria) + cut_tight;
    TH1D *h_ttbar = load_histogram("mc09_ttbar", prefix_hist, observable, cut, nbins, xlow, xhigh);
    h_ttbar->Scale( getIntLumi(sample) );

    // stop
    cut = mycut("mc09_stop", jetNumberCriteria) + cut_tight;
    TH1D *h_stop = load_histogram("mc09_stop", prefix_hist, observable, cut, nbins, xlow, xhigh);
    h_stop->Scale( getIntLumi(sample) );

    // wjets
    cut = mycut("mc09_wjets", jetNumberCriteria) + cut_tight;
    TH1D *h_wjets = load_histogram("mc09_wjets", prefix_hist, observable, cut, nbins, xlow, xhigh);
    h_wjets->Scale( getIntLumi(sample) );

    // zjets
    cut = mycut("mc09_zjets", jetNumberCriteria) + cut_tight;
    TH1D *h_zjets = load_histogram("mc09_zjets", prefix_hist, observable, cut, nbins, xlow, xhigh);
    h_zjets->Scale( getIntLumi(sample) );

    // diboson 
    cut = mycut("mc09_diboson", jetNumberCriteria) + cut_tight;
    TH1D *h_diboson = load_histogram("mc09_diboson", prefix_hist, observable, cut, nbins, xlow, xhigh);
    h_diboson->Scale( getIntLumi(sample) );

    // data-driven QCD estimation
    TH1D *h_qcd = new TH1D("h_qcd", "h_qcd", nbins, xlow, xhigh);

    cut = mycut(sample, jetNumberCriteria) + cut_loose;
    TH1D *h_loose = load_histogram(sample, "h_qcd_loose_", observable, cut, nbins, xlow, xhigh);
    h_qcd = PlotQCD(h_loose, h_data, 4, 0, 2);
    h_qcd->SetNameTitle("h_qcd","h_qcd");

    // writing out the histograms 
    TFile *rsult = new TFile(prefix_hist + "output.root", "recreate");

    h_data->Write();
    h_ttbar->Write();
    h_stop->Write();
    h_wjets->Write();
    h_zjets->Write();
    h_diboson->Write();

    h_qcd->Write();

    rsult->Close();
 
    gDirectory->Clear();

}

void plot_MWt(const char* sample="data10", const char* jetNumberCriteria=">=1") {

    TString cut;

    TString prefix_hist = "h_";

    TString cut_loose = " && GoodMuLoose_dR_jet > 0.4 ";
    TString cut_tight = cut_loose + " && GoodMuLoose_etcone30 < 4000. && GoodMuLoose_ptcone30 < 4000. ";
    TString eq_mwt = "sqrt( pow( topMET_et + GoodMuLoose_pt, 2) - pow( topMET_etx + GoodMuLoose_px, 2) - pow( topMET_ety + GoodMuLoose_py,2) )";

    // data
    cut = mycut(sample, jetNumberCriteria) + cut_tight;
    TH1D *h_data = load_histogram(sample, prefix_hist, eq_mwt, cut, 30, 0., 150000.);
    h_data->SetNameTitle(prefix_hist + "data", prefix_hist + "data");

    // ttbar
    cut = mycut("mc09_ttbar", jetNumberCriteria) + " && MWtrans > 0.";
    TH1D *h_ttbar = load_histogram("mc09_ttbar", prefix_hist, "MWtrans", cut, 30, 0., 150000.);
    h_ttbar->Scale( getIntLumi(sample) );

    // stop
    cut = mycut("mc09_stop", jetNumberCriteria) + " && MWtrans > 0.";
    TH1D *h_stop = load_histogram("mc09_stop", prefix_hist, "MWtrans", cut, 30, 0., 150000.);
    h_stop->Scale( getIntLumi(sample) );

    // wjets
    cut = mycut("mc09_wjets", jetNumberCriteria) + " && MWtrans > 0.";
    TH1D *h_wjets = load_histogram("mc09_wjets", prefix_hist, "MWtrans", cut, 30, 0., 150000.);
    h_wjets->Scale( getIntLumi(sample) );

    // zjets
    cut = mycut("mc09_zjets", jetNumberCriteria) + " && MWtrans > 0.";
    TH1D *h_zjets = load_histogram("mc09_zjets", prefix_hist, "MWtrans", cut, 30, 0., 150000.);
    h_zjets->Scale( getIntLumi(sample) );

    // diboson 
    cut = mycut("mc09_diboson", jetNumberCriteria) + " && MWtrans > 0.";
    TH1D *h_diboson = load_histogram("mc09_diboson", prefix_hist, "MWtrans", cut, 30, 0., 150000.);
    h_diboson->Scale( getIntLumi(sample) );

    // data-driven QCD estimation
    TH1D *h_qcd = new TH1D("h_qcd", "h_qcd", 30, 0., 150000.);

    cut = mycut(sample, jetNumberCriteria) + cut_loose;
    TH1D *h_loose = load_histogram(sample, "h_qcd_loose_", eq_mwt, cut, 30, 0., 150000.);
    h_qcd = PlotQCD(h_loose, h_data, 4, 0, 2);
    h_qcd->SetNameTitle("h_qcd","h_qcd");

    // writing out the h_eff
    TFile *rsult = new TFile(prefix_hist + "output.root", "recreate");

    h_data->Write();
    h_ttbar->Write();
    h_stop->Write();
    h_wjets->Write();
    h_zjets->Write();
    h_diboson->Write();

    h_qcd->Write();

    rsult->Close();
 
    gDirectory->Clear();
}

void plot_topMass(const char* sample="data10", const char* jetNumberCriteria=">=1") {
    plot_observable("topMass_GoodJet", "h_", 16, 50000, 450000, sample, jetNumberCriteria); 
}
